from pycocotools.coco import COCO
import matplotlib.pyplot as plt
import logging
import cv2
from tqdm import tqdm
import os
import numpy as np
import random
from pathlib import Path

COCOROOT = Path("./dataset/arcade_challenge_datasets/dataset_phase_1/segmentation_dataset/seg_val")
ANNFILE = list(COCOROOT.glob("./annotations/" + "*.json"))[0]



def coco2mask(coco_inst:COCO, idx:int):
    img_path = os.path.join(COCOROOT, "images", coco_inst.imgs[idx]["file_name"])
    
    annIds = coco_inst.getAnnIds(imgIds=idx, iscrowd=None)
    anns = coco_inst.loadAnns(annIds)
    output_mask = None
    for ann in anns:
        category_id = ann['category_id']
        temp_mask = coco_inst.annToMask(ann) 
        if output_mask is None: output_mask = temp_mask * category_id
        else: output_mask += temp_mask * category_id
    return output_mask

def main():
    coco=COCO(ANNFILE)
    print(coco.loadCats(coco.getCatIds()))
    os.makedirs(os.path.join(COCOROOT, "masks"), exist_ok=True)
    mask_dir = os.path.join(COCOROOT, "masks")

    for idx in tqdm(coco.getImgIds()):
        msk = coco2mask(coco, idx)
        f_name = coco.imgs[idx]["file_name"]
        out_path = os.path.join(mask_dir, f_name)
        cv2.imwrite(out_path, msk)

if __name__ == '__main__':
    main()